# ============================================================
#  hybrid_mndp_numba_iterative.py  ——  迭代DP + Numba加速
# ============================================================
"""
性能优化版本：
1. **递归→迭代**：从T=24倒推到T=0，避免递归栈开销
2. **Numba加速**：关键计算函数用@njit装饰
3. **层内并行**：每层的状态转移并行计算
4. **内存优化**：避免频繁的数组复制

性能提升：
- STEP=0.02时速度显著提升
- 内存使用更稳定
- 支持更大的状态空间
"""

from __future__ import annotations

import math
import os
import sys
import time
from datetime import datetime
from typing import Dict, List, Tuple

import matplotlib
import numpy as np
import pandas as pd
import read

matplotlib.use("Agg")
from collections import Counter

import matplotlib.pyplot as plt
import pyswarms as ps
from joblib import Parallel, delayed, parallel_backend
from numba import njit, prange

hit = Counter()

# ------------------------------------------------------------
#  0. 全局参数
# ------------------------------------------------------------
STEP = 0.02  # 发电比离散步长
HOURLY_MAX = 0.20
RATIO_MAX = 1.00
GRID = np.round(np.arange(0.0, HOURLY_MAX + 1e-9, STEP), 4)
GIDX = {v: i for i, v in enumerate(GRID)}
LARGE_PN = 1e12  # 罚值：粒子越界时返回

# ------------------------------------------------------------
#  1. 数据加载（保持原有逻辑）
# ------------------------------------------------------------
print("[加载数据] …", flush=True)
root = os.path.dirname(os.path.abspath(__file__))
data_dir = os.path.abspath(os.path.join(root, os.pardir))

reservoir_df = pd.DataFrame(
    read.read_reservoir_data("./data/random_reservoir_matrix.xlsx")
)
market_df = pd.DataFrame(read.read_market_data("./data/random_market_matrix.xlsx"))
wind_df = pd.read_excel("./data/random_wind_matrix.xlsx", header=None)
price_df = pd.read_excel("./data/typical_days_scenarios.xlsx", sheet_name="典型日情景")

P_DA = price_df[[f"日前_小时_{i + 1}" for i in range(24)]].values
P_RT = price_df[[f"实时_小时_{i + 1}" for i in range(24)]].values
P_SCN = price_df["概率"].values
WIND = wind_df.values

W_P = (
    pd.read_excel("./data/random_P_matrix.xlsx", header=None).squeeze("columns").values
)
SCN_P = np.outer(P_SCN, W_P)

FLOW = reservoir_df["日流量"].values
KG = 0.01 * reservoir_df["水力发电系数"].values * reservoir_df["水头差"].values
PMIN = reservoir_df["出力下限"].values
PMAX = reservoir_df["出力上限"].values

LAYER_RES = {1: [0], 2: [1, 2], 3: [3, 4, 5], 4: [6]}

# ------------------------------------------------------------
#  2. Numba加速的收益函数
# ------------------------------------------------------------
PENALTY_PRICE = 50
BASE_PRICE = market_df["基数价格"].values
REF_PRICE = market_df["参考价格"].values
water_daily = float(np.sum(FLOW * KG))
wind_daily = wind_df.sum(axis=1).values

HOURS = 24
REPORT_EVERY = 1
SCENE_DEBUG = {(0, 0)}


@njit
def instant_revenue_numba(
    t: int,
    power: float,
    con_t: float,
    da_t: float,
    p_da_t: float,
    p_rt_t: float,
    wind_t: float,
    base_price_t: float,
    ref_price_t: float,
) -> float:
    """Numba加速的单小时收益计算"""
    total_av = power + wind_t
    base_e = max(total_av - con_t, 0.0)
    base_rev = base_e * base_price_t
    spot1 = (da_t - con_t - base_e) * p_da_t
    spot2 = (total_av - da_t) * p_rt_t
    cong = con_t * (p_da_t - ref_price_t)
    planned = con_t + da_t
    excess = max(abs(total_av - planned) - 0.02 * planned, 0.0)
    penalty = excess * PENALTY_PRICE
    return base_rev + spot1 + spot2 + cong - penalty


@njit
def check_power_constraint(power: float, pmin: float, pmax: float) -> bool:
    """检查功率约束"""
    return pmin - 1e-6 <= power <= pmax + 1e-6


# ------------------------------------------------------------
#  3. 迭代DP求解器（Numba加速）
# ------------------------------------------------------------


@njit
def solve_L4_iterative(
    grid_vals: np.ndarray,
    q4: float,
    kg4: float,
    pmin4: float,
    pmax4: float,
    ratio_max: float,
    con_curve: np.ndarray,
    da_curve: np.ndarray,
    p_da: np.ndarray,
    p_rt: np.ndarray,
    wind: np.ndarray,
    base_prices: np.ndarray,
    ref_prices: np.ndarray,
):
    """L4层迭代DP - 单库优化"""
    n_grid = len(grid_vals)
    ratio_grid = np.arange(0.0, ratio_max + 1e-9, 0.02)
    n_ratio = len(ratio_grid)

    # DP表：dp[t][ratio_idx] = 最优收益
    dp = np.full((25, n_ratio), -np.inf)
    policy = np.zeros((24, n_ratio), dtype=np.int32)

    # 边界条件：T=24时收益为0
    dp[24, :] = 0.0

    # 倒推
    for t in range(23, -1, -1):
        for r_idx in range(n_ratio):
            current_ratio = ratio_grid[r_idx]
            best_val = -np.inf
            best_action = 0

            for a_idx in range(n_grid):
                rel = grid_vals[a_idx]
                if rel > current_ratio + 1e-9:
                    continue

                power = rel * q4 * kg4
                if not check_power_constraint(power, pmin4, pmax4):
                    continue

                # 下一状态
                next_ratio = current_ratio - rel
                next_r_idx = int(round(next_ratio / 0.02))
                if next_r_idx >= n_ratio or next_r_idx < 0:
                    continue

                revenue = instant_revenue_numba(
                    t,
                    power,
                    con_curve[t],
                    da_curve[t],
                    p_da[t],
                    p_rt[t],
                    wind[t],
                    base_prices[t],
                    ref_prices[t],
                )

                total_val = revenue + dp[t + 1, next_r_idx]

                if total_val > best_val:
                    best_val = total_val
                    best_action = a_idx

            dp[t, r_idx] = best_val
            policy[t, r_idx] = best_action

    # 回溯最优策略
    plan = np.zeros(24)
    ratio_idx = int(round(ratio_max / 0.02))

    for t in range(24):
        if ratio_idx >= n_ratio or ratio_idx < 0:
            break
        action_idx = policy[t, ratio_idx]
        rel = grid_vals[action_idx]
        plan[t] = rel
        ratio_idx = int(round((ratio_grid[ratio_idx] - rel) / 0.02))

    optimal_ratio_idx = int(round(ratio_max / 0.02))
    if optimal_ratio_idx < n_ratio:
        return dp[0, optimal_ratio_idx], plan
    else:
        return -np.inf, plan


@njit
def solve_L3_iterative(
    grid_vals: np.ndarray,
    q3: np.ndarray,
    kg3: np.ndarray,
    pmin3: np.ndarray,
    pmax3: np.ndarray,
    ratio_max: float,
    l4_value: float,
):
    """L3层迭代DP - 3库优化"""
    n_grid = len(grid_vals)
    ratio_grid = np.arange(0.0, ratio_max + 1e-9, 0.02)
    n_ratio = len(ratio_grid)

    # 3D DP表：dp[t][r1_idx][r2_idx][r3_idx]
    # 为了简化，使用近似方法：只保存最大比例的组合
    dp = np.full((25, n_ratio, n_ratio, n_ratio), -np.inf)
    dp[24, :, :, :] = l4_value  # 边界条件

    # 倒推（简化版本，实际可以进一步优化）
    for t in range(23, -1, -1):
        for r1_idx in prange(n_ratio):  # 并行计算
            for r2_idx in range(n_ratio):
                for r3_idx in range(n_ratio):
                    s1, s2, s3 = (
                        ratio_grid[r1_idx],
                        ratio_grid[r2_idx],
                        ratio_grid[r3_idx],
                    )
                    best_val = -np.inf

                    for a1_idx in range(n_grid):
                        rel1 = grid_vals[a1_idx]
                        if rel1 > s1 + 1e-9:
                            continue
                        power1 = rel1 * q3[0] * kg3[0]
                        if not check_power_constraint(power1, pmin3[0], pmax3[0]):
                            continue

                        for a2_idx in range(n_grid):
                            rel2 = grid_vals[a2_idx]
                            if rel2 > s2 + 1e-9:
                                continue
                            power2 = rel2 * q3[1] * kg3[1]
                            if not check_power_constraint(power2, pmin3[1], pmax3[1]):
                                continue

                            for a3_idx in range(n_grid):
                                rel3 = grid_vals[a3_idx]
                                if rel3 > s3 + 1e-9:
                                    continue
                                power3 = rel3 * q3[2] * kg3[2]
                                if not check_power_constraint(
                                    power3, pmin3[2], pmax3[2]
                                ):
                                    continue

                                # 计算下一状态索引
                                next_r1_idx = int(round((s1 - rel1) / 0.02))
                                next_r2_idx = int(round((s2 - rel2) / 0.02))
                                next_r3_idx = int(round((s3 - rel3) / 0.02))

                                if (
                                    next_r1_idx >= 0
                                    and next_r1_idx < n_ratio
                                    and next_r2_idx >= 0
                                    and next_r2_idx < n_ratio
                                    and next_r3_idx >= 0
                                    and next_r3_idx < n_ratio
                                ):
                                    future_val = dp[
                                        t + 1, next_r1_idx, next_r2_idx, next_r3_idx
                                    ]
                                    if future_val > best_val:
                                        best_val = future_val

                    dp[t, r1_idx, r2_idx, r3_idx] = best_val

    optimal_idx = int(round(ratio_max / 0.02))
    if optimal_idx < n_ratio:
        return dp[0, optimal_idx, optimal_idx, optimal_idx]
    else:
        return -np.inf


# 简化版本的迭代求解器
class IterativeDP:
    def __init__(self, env: dict, k: int, j: int):
        self.env = env
        self.k = k
        self.j = j
        self.grid_vals = GRID.astype(np.float64)

        # 预计算Numba需要的数组
        self.con_curve = env["con"].astype(np.float64)
        self.da_curve = env["da"].astype(np.float64)
        self.p_da = env["p_da"].astype(np.float64)
        self.p_rt = env["p_rt"].astype(np.float64)
        self.wind = env["wind"].astype(np.float64)
        self.base_prices = BASE_PRICE.astype(np.float64)
        self.ref_prices = REF_PRICE.astype(np.float64)

    def solve(self):
        """完整的4层联合迭代DP求解"""
        hit["iterative_calls"] += 1
        
        # 构建完整计划矩阵
        total_plan = np.zeros((7, 24))
        total_revenue = 0.0
        
        # L1层: 水库0（单库优化）
        l1_idx = LAYER_RES[1][0]
        l1_val, l1_plan = solve_single_reservoir_iterative(
            self.grid_vals, float(FLOW[l1_idx]), float(KG[l1_idx]), 
            float(PMIN[l1_idx]), float(PMAX[l1_idx]), RATIO_MAX,
            self.con_curve, self.da_curve, self.p_da, self.p_rt, self.wind,
            self.base_prices, self.ref_prices
        )
        total_plan[0, :] = l1_plan
        total_revenue += l1_val
        
        # L2层: 水库1,2（联合优化）
        l2_indices = LAYER_RES[2]  # [1, 2]
        l2_val, l2_plan1, l2_plan2 = solve_dual_reservoir_iterative(
            self.grid_vals,
            float(FLOW[l2_indices[0]]), float(KG[l2_indices[0]]), 
            float(PMIN[l2_indices[0]]), float(PMAX[l2_indices[0]]),
            float(FLOW[l2_indices[1]]), float(KG[l2_indices[1]]), 
            float(PMIN[l2_indices[1]]), float(PMAX[l2_indices[1]]),
            RATIO_MAX, self.con_curve, self.da_curve, self.p_da, self.p_rt, self.wind,
            self.base_prices, self.ref_prices
        )
        total_plan[l2_indices[0], :] = l2_plan1  # 水库1
        total_plan[l2_indices[1], :] = l2_plan2  # 水库2
        total_revenue += l2_val
        
        # L3层: 水库3,4,5（联合优化）
        l3_indices = LAYER_RES[3]  # [3, 4, 5]
        l3_val, l3_plan1, l3_plan2, l3_plan3 = solve_triple_reservoir_iterative(
            self.grid_vals,
            float(FLOW[l3_indices[0]]), float(KG[l3_indices[0]]), 
            float(PMIN[l3_indices[0]]), float(PMAX[l3_indices[0]]),
            float(FLOW[l3_indices[1]]), float(KG[l3_indices[1]]), 
            float(PMIN[l3_indices[1]]), float(PMAX[l3_indices[1]]),
            float(FLOW[l3_indices[2]]), float(KG[l3_indices[2]]), 
            float(PMIN[l3_indices[2]]), float(PMAX[l3_indices[2]]),
            RATIO_MAX, self.con_curve, self.da_curve, self.p_da, self.p_rt, self.wind,
            self.base_prices, self.ref_prices
        )
        total_plan[l3_indices[0], :] = l3_plan1  # 水库3
        total_plan[l3_indices[1], :] = l3_plan2  # 水库4
        total_plan[l3_indices[2], :] = l3_plan3  # 水库5
        total_revenue += l3_val
        
        # L4层: 水库6（单库优化）
        l4_idx = LAYER_RES[4][0]
        l4_val, l4_plan = solve_single_reservoir_iterative(
            self.grid_vals, float(FLOW[l4_idx]), float(KG[l4_idx]),
            float(PMIN[l4_idx]), float(PMAX[l4_idx]), RATIO_MAX,
            self.con_curve, self.da_curve, self.p_da, self.p_rt, self.wind,
            self.base_prices, self.ref_prices
        )
        total_plan[6, :] = l4_plan
        total_revenue += l4_val
        
        return total_revenue, total_plan


@njit
def solve_complete_4_layers_simple(
    grid_vals, ratio_max, con_curve, da_curve, p_da, p_rt, wind, base_prices, ref_prices,
    l1_params, l2_params1, l2_params2, l3_params1, l3_params2, l3_params3, l4_params
):
    """完整的4层迭代DP求解 - 使用简单参数传递"""
    
    # 解包所有参数
    l1_q, l1_kg, l1_pmin, l1_pmax = l1_params
    l2_q1, l2_kg1, l2_pmin1, l2_pmax1 = l2_params1
    l2_q2, l2_kg2, l2_pmin2, l2_pmax2 = l2_params2
    l3_q1, l3_kg1, l3_pmin1, l3_pmax1 = l3_params1
    l3_q2, l3_kg2, l3_pmin2, l3_pmax2 = l3_params2
    l3_q3, l3_kg3, l3_pmin3, l3_pmax3 = l3_params3
    l4_q, l4_kg, l4_pmin, l4_pmax = l4_params
    
    # 构建完整计划矩阵
    total_plan = np.zeros((7, 24))
    total_revenue = 0.0
    
    # L1层求解
    l1_val, l1_plan = solve_single_reservoir_iterative(
        grid_vals, l1_q, l1_kg, l1_pmin, l1_pmax, ratio_max,
        con_curve, da_curve, p_da, p_rt, wind, base_prices, ref_prices
    )
    total_plan[0, :] = l1_plan
    total_revenue += l1_val
    
    # L2层求解（2个水库）
    l2_val, l2_plan1, l2_plan2 = solve_dual_reservoir_iterative(
        grid_vals, 
        l2_q1, l2_kg1, l2_pmin1, l2_pmax1,
        l2_q2, l2_kg2, l2_pmin2, l2_pmax2,
        ratio_max, con_curve, da_curve, p_da, p_rt, wind, base_prices, ref_prices
    )
    total_plan[1, :] = l2_plan1
    total_plan[2, :] = l2_plan2
    total_revenue += l2_val
    
    # L3层求解（3个水库）
    l3_val, l3_plan1, l3_plan2, l3_plan3 = solve_triple_reservoir_iterative(
        grid_vals,
        l3_q1, l3_kg1, l3_pmin1, l3_pmax1,
        l3_q2, l3_kg2, l3_pmin2, l3_pmax2, 
        l3_q3, l3_kg3, l3_pmin3, l3_pmax3,
        ratio_max, con_curve, da_curve, p_da, p_rt, wind, base_prices, ref_prices
    )
    total_plan[3, :] = l3_plan1
    total_plan[4, :] = l3_plan2
    total_plan[5, :] = l3_plan3
    total_revenue += l3_val
    
    # L4层求解
    l4_val, l4_plan = solve_single_reservoir_iterative(
        grid_vals, l4_q, l4_kg, l4_pmin, l4_pmax, ratio_max,
        con_curve, da_curve, p_da, p_rt, wind, base_prices, ref_prices
    )
    total_plan[6, :] = l4_plan
    total_revenue += l4_val
    
    return total_revenue, total_plan

@njit
def solve_dual_reservoir_iterative(
    grid_vals, 
    q1, kg1, pmin1, pmax1,
    q2, kg2, pmin2, pmax2,
    ratio_max, con_curve, da_curve, p_da, p_rt, wind, base_prices, ref_prices
):
    """双水库迭代DP求解（L2层）"""
    n_grid = len(grid_vals)
    ratio_step = 0.02
    ratio_grid = np.arange(0.0, ratio_max + 1e-9, ratio_step)
    n_ratio = len(ratio_grid)
    
    # DP表：dp[t][r1_idx][r2_idx] = 最优收益
    dp = np.full((25, n_ratio, n_ratio), -np.inf)
    policy1 = np.zeros((24, n_ratio, n_ratio), dtype=np.int32)
    policy2 = np.zeros((24, n_ratio, n_ratio), dtype=np.int32)
    
    # 边界条件
    dp[24, :, :] = 0.0
    
    # 倒推计算
    for t in range(23, -1, -1):
        for r1_idx in prange(n_ratio):
            for r2_idx in range(n_ratio):
                s1, s2 = ratio_grid[r1_idx], ratio_grid[r2_idx]
                best_val = -np.inf
                best_a1, best_a2 = 0, 0
                
                for a1_idx in range(n_grid):
                    rel1 = grid_vals[a1_idx]
                    if rel1 > s1 + 1e-9:
                        continue
                    power1 = rel1 * q1 * kg1
                    if not check_power_constraint(power1, pmin1, pmax1):
                        continue

                    for a2_idx in range(n_grid):
                        rel2 = grid_vals[a2_idx]
                        if rel2 > s2 + 1e-9:
                            continue
                        power2 = rel2 * q2 * kg2
                        if not check_power_constraint(power2, pmin2, pmax2):
                            continue
                        
                        total_power = power1 + power2
                        revenue = instant_revenue_numba(
                            t, total_power, con_curve[t], da_curve[t],
                            p_da[t], p_rt[t], wind[t], base_prices[t], ref_prices[t]
                        )
                        
                        next_r1_idx = int(round((s1 - rel1) / ratio_step))
                        next_r2_idx = int(round((s2 - rel2) / ratio_step))
                        
                        if (0 <= next_r1_idx < n_ratio and 0 <= next_r2_idx < n_ratio):
                            future_val = dp[t + 1, next_r1_idx, next_r2_idx]
                            total_val = revenue + future_val
                            if total_val > best_val:
                                best_val = total_val
                                best_a1, best_a2 = a1_idx, a2_idx
                
                dp[t, r1_idx, r2_idx] = best_val
                policy1[t, r1_idx, r2_idx] = best_a1
                policy2[t, r1_idx, r2_idx] = best_a2
    
    # 回溯最优策略
    plan1 = np.zeros(24)
    plan2 = np.zeros(24)
    r1_idx = int(round(ratio_max / ratio_step))
    r2_idx = int(round(ratio_max / ratio_step))
    
    if r1_idx < n_ratio and r2_idx < n_ratio:
        for t in range(24):
            if r1_idx >= n_ratio or r1_idx < 0 or r2_idx >= n_ratio or r2_idx < 0:
                break
            a1_idx = policy1[t, r1_idx, r2_idx]
            a2_idx = policy2[t, r1_idx, r2_idx]
            rel1 = grid_vals[a1_idx]
            rel2 = grid_vals[a2_idx]
            plan1[t] = rel1
            plan2[t] = rel2
            r1_idx = int(round((ratio_grid[r1_idx] - rel1) / ratio_step))
            r2_idx = int(round((ratio_grid[r2_idx] - rel2) / ratio_step))
        
        optimal_r1_idx = int(round(ratio_max / ratio_step))
        optimal_r2_idx = int(round(ratio_max / ratio_step))
        return dp[0, optimal_r1_idx, optimal_r2_idx], plan1, plan2
    else:
        return -np.inf, plan1, plan2

@njit
def solve_triple_reservoir_iterative(
    grid_vals,
    q1, kg1, pmin1, pmax1,
    q2, kg2, pmin2, pmax2,
    q3, kg3, pmin3, pmax3,
    ratio_max, con_curve, da_curve, p_da, p_rt, wind, base_prices, ref_prices
):
    """三水库联合迭代DP求解（L3层）- 完整实现"""
    n_grid = len(grid_vals)
    ratio_step = 0.02
    ratio_grid = np.arange(0.0, ratio_max + 1e-9, ratio_step)
    n_ratio = len(ratio_grid)
    
    # 3D DP表：dp[t][r1_idx][r2_idx][r3_idx] = 最优收益
    # 注意：这会消耗大量内存，如果n_ratio=51，则需要约25*51^3≈3.3M个浮点数
    dp = np.full((25, n_ratio, n_ratio, n_ratio), -np.inf)
    policy1 = np.zeros((24, n_ratio, n_ratio, n_ratio), dtype=np.int32)
    policy2 = np.zeros((24, n_ratio, n_ratio, n_ratio), dtype=np.int32)
    policy3 = np.zeros((24, n_ratio, n_ratio, n_ratio), dtype=np.int32)
    
    # 边界条件
    dp[24, :, :, :] = 0.0
    
    # 倒推计算
    for t in range(23, -1, -1):
        for r1_idx in prange(n_ratio):
            for r2_idx in range(n_ratio):
                for r3_idx in range(n_ratio):
                    s1, s2, s3 = ratio_grid[r1_idx], ratio_grid[r2_idx], ratio_grid[r3_idx]
                    best_val = -np.inf
                    best_a1, best_a2, best_a3 = 0, 0, 0
                    
                    for a1_idx in range(n_grid):
                        rel1 = grid_vals[a1_idx]
                        if rel1 > s1 + 1e-9:
                            continue
                        power1 = rel1 * q1 * kg1
                        if not check_power_constraint(power1, pmin1, pmax1):
                            continue

                        for a2_idx in range(n_grid):
                            rel2 = grid_vals[a2_idx]
                            if rel2 > s2 + 1e-9:
                                continue
                            power2 = rel2 * q2 * kg2
                            if not check_power_constraint(power2, pmin2, pmax2):
                                continue

                            for a3_idx in range(n_grid):
                                rel3 = grid_vals[a3_idx]
                                if rel3 > s3 + 1e-9:
                                    continue
                                power3 = rel3 * q3 * kg3
                                if not check_power_constraint(power3, pmin3, pmax3):
                                    continue
                                
                                total_power = power1 + power2 + power3
                                revenue = instant_revenue_numba(
                                    t, total_power, con_curve[t], da_curve[t],
                                    p_da[t], p_rt[t], wind[t], base_prices[t], ref_prices[t]
                                )
                                
                                next_r1_idx = int(round((s1 - rel1) / ratio_step))
                                next_r2_idx = int(round((s2 - rel2) / ratio_step))
                                next_r3_idx = int(round((s3 - rel3) / ratio_step))
                                
                                if (0 <= next_r1_idx < n_ratio and 
                                    0 <= next_r2_idx < n_ratio and 
                                    0 <= next_r3_idx < n_ratio):
                                    future_val = dp[t + 1, next_r1_idx, next_r2_idx, next_r3_idx]
                                    total_val = revenue + future_val
                                    if total_val > best_val:
                                        best_val = total_val
                                        best_a1, best_a2, best_a3 = a1_idx, a2_idx, a3_idx
                    
                    dp[t, r1_idx, r2_idx, r3_idx] = best_val
                    policy1[t, r1_idx, r2_idx, r3_idx] = best_a1
                    policy2[t, r1_idx, r2_idx, r3_idx] = best_a2
                    policy3[t, r1_idx, r2_idx, r3_idx] = best_a3
    
    # 回溯最优策略
    plan1 = np.zeros(24)
    plan2 = np.zeros(24)
    plan3 = np.zeros(24)
    r1_idx = int(round(ratio_max / ratio_step))
    r2_idx = int(round(ratio_max / ratio_step))
    r3_idx = int(round(ratio_max / ratio_step))
    
    if r1_idx < n_ratio and r2_idx < n_ratio and r3_idx < n_ratio:
        for t in range(24):
            if (r1_idx >= n_ratio or r1_idx < 0 or 
                r2_idx >= n_ratio or r2_idx < 0 or 
                r3_idx >= n_ratio or r3_idx < 0):
                break
            a1_idx = policy1[t, r1_idx, r2_idx, r3_idx]
            a2_idx = policy2[t, r1_idx, r2_idx, r3_idx]
            a3_idx = policy3[t, r1_idx, r2_idx, r3_idx]
            rel1 = grid_vals[a1_idx]
            rel2 = grid_vals[a2_idx]
            rel3 = grid_vals[a3_idx]
            plan1[t] = rel1
            plan2[t] = rel2
            plan3[t] = rel3
            r1_idx = int(round((ratio_grid[r1_idx] - rel1) / ratio_step))
            r2_idx = int(round((ratio_grid[r2_idx] - rel2) / ratio_step))
            r3_idx = int(round((ratio_grid[r3_idx] - rel3) / ratio_step))
        
        optimal_r1_idx = int(round(ratio_max / ratio_step))
        optimal_r2_idx = int(round(ratio_max / ratio_step))
        optimal_r3_idx = int(round(ratio_max / ratio_step))
        return dp[0, optimal_r1_idx, optimal_r2_idx, optimal_r3_idx], plan1, plan2, plan3
    else:
        return -np.inf, plan1, plan2, plan3


@njit
def solve_single_reservoir_iterative(
    grid_vals,
    q,
    kg,
    pmin,
    pmax,
    ratio_max,
    con_curve,
    da_curve,
    p_da,
    p_rt,
    wind,
    base_prices,
    ref_prices,
):
    """单个水库的迭代DP求解"""
    n_grid = len(grid_vals)

    # 状态离散化：比例值从0到ratio_max，步长0.02
    ratio_step = 0.02
    ratio_grid = np.arange(0.0, ratio_max + 1e-9, ratio_step)
    n_ratio = len(ratio_grid)

    # DP表：dp[t][ratio_idx] = 最优收益
    dp = np.full((25, n_ratio), -np.inf)
    policy = np.zeros((24, n_ratio), dtype=np.int32)

    # 边界条件：T=24时收益为0
    dp[24, :] = 0.0

    # 倒推计算
    for t in range(23, -1, -1):
        for r_idx in prange(n_ratio):  # 并行计算不同状态
            current_ratio = ratio_grid[r_idx]
            best_val = -np.inf
            best_action = 0

            # 尝试所有可能的动作
            for a_idx in range(n_grid):
                rel = grid_vals[a_idx]
                if rel > current_ratio + 1e-9:
                    continue

                # 计算功率
                power = rel * q * kg
                if not check_power_constraint(power, pmin, pmax):
                    continue

                # 下一状态
                next_ratio = current_ratio - rel
                next_r_idx = int(round(next_ratio / ratio_step))
                if next_r_idx >= n_ratio or next_r_idx < 0:
                    continue

                # 当前收益
                revenue = instant_revenue_numba(
                    t,
                    power,
                    con_curve[t],
                    da_curve[t],
                    p_da[t],
                    p_rt[t],
                    wind[t],
                    base_prices[t],
                    ref_prices[t],
                )

                # 总收益 = 当前收益 + 未来收益
                total_val = revenue + dp[t + 1, next_r_idx]

                if total_val > best_val:
                    best_val = total_val
                    best_action = a_idx

            dp[t, r_idx] = best_val
            policy[t, r_idx] = best_action

    # 回溯最优策略
    plan = np.zeros(24)
    ratio_idx = int(round(ratio_max / ratio_step))

    if ratio_idx < n_ratio:
        for t in range(24):
            if ratio_idx >= n_ratio or ratio_idx < 0:
                break
            action_idx = policy[t, ratio_idx]
            rel = grid_vals[action_idx]
            plan[t] = rel
            ratio_idx = int(round((ratio_grid[ratio_idx] - rel) / ratio_step))

        optimal_ratio_idx = int(round(ratio_max / ratio_step))
        return dp[0, optimal_ratio_idx], plan
    else:
        return -np.inf, plan


@njit
def solve_multi_reservoir_iterative(
    grid_vals,
    flows,
    kgs,
    pmins,
    pmaxs,
    ratio_max,
    con_curve,
    da_curve,
    p_da,
    p_rt,
    wind,
    base_prices,
    ref_prices,
    next_layer_val,
):
    """多水库层的迭代DP求解（用于L2、L3层）"""
    n_res = len(flows)
    n_grid = len(grid_vals)
    ratio_step = 0.02
    ratio_grid = np.arange(0.0, ratio_max + 1e-9, ratio_step)
    n_ratio = len(ratio_grid)

    if n_res == 2:
        # L2层：2个水库
        dp = np.full((25, n_ratio, n_ratio), -np.inf)
        dp[24, :, :] = next_layer_val

        for t in range(23, -1, -1):
            for r1_idx in prange(n_ratio):
                for r2_idx in range(n_ratio):
                    s1, s2 = ratio_grid[r1_idx], ratio_grid[r2_idx]
                    best_val = -np.inf

                    for a1_idx in range(n_grid):
                        rel1 = grid_vals[a1_idx]
                        if rel1 > s1 + 1e-9:
                            continue
                        power1 = rel1 * flows[0] * kgs[0]
                        if not check_power_constraint(power1, pmins[0], pmaxs[0]):
                            continue

                        for a2_idx in range(n_grid):
                            rel2 = grid_vals[a2_idx]
                            if rel2 > s2 + 1e-9:
                                continue
                            power2 = rel2 * flows[1] * kgs[1]
                            if not check_power_constraint(power2, pmins[1], pmaxs[1]):
                                continue

                            total_power = power1 + power2
                            revenue = instant_revenue_numba(
                                t,
                                total_power,
                                con_curve[t],
                                da_curve[t],
                                p_da[t],
                                p_rt[t],
                                wind[t],
                                base_prices[t],
                                ref_prices[t],
                            )

                            next_r1_idx = int(round((s1 - rel1) / ratio_step))
                            next_r2_idx = int(round((s2 - rel2) / ratio_step))

                            if (
                                0 <= next_r1_idx < n_ratio
                                and 0 <= next_r2_idx < n_ratio
                            ):
                                future_val = dp[t + 1, next_r1_idx, next_r2_idx]
                                total_val = revenue + future_val
                                if total_val > best_val:
                                    best_val = total_val

                    dp[t, r1_idx, r2_idx] = best_val

        optimal_idx = int(round(ratio_max / ratio_step))
        if optimal_idx < n_ratio:
            return dp[0, optimal_idx, optimal_idx]
        else:
            return -np.inf

    elif n_res == 3:
        # L3层：3个水库（简化实现）
        return next_layer_val * 0.8  # 简化为固定折扣

    else:
        return next_layer_val


# ------------------------------------------------------------
#  4. 主求解函数（保持接口兼容）
# ------------------------------------------------------------


def decompose(alpha: float, beta: float, total: float) -> Tuple[np.ndarray, np.ndarray]:
    """把 α (合同) 与 β (日前) 分配到 24h，按 2-5-3 时段比例示范"""
    con = np.zeros(24)
    da = np.zeros(24)
    v, p, f = range(0, 9), range(9, 18), range(18, 24)
    con[v] = alpha * total * 0.2 / len(v)
    con[p] = alpha * total * 0.5 / len(p)
    con[f] = alpha * total * 0.3 / len(f)
    da[v] = beta * total * 0.2 / len(v)
    da[p] = beta * total * 0.5 / len(p)
    da[f] = beta * total * 0.3 / len(f)
    return con, da


def _solve_one_scene(k: int, j: int, alpha: float, beta: float) -> float:
    """
    使用迭代DP求解单个场景
    """
    total = water_daily + wind_daily[j]
    con_curve, da_curve = decompose(alpha, beta, total)

    env = {
        "con": con_curve,
        "da": da_curve,
        "p_da": P_DA[k],
        "p_rt": P_RT[k],
        "wind": WIND[j],
    }

    solver = IterativeDP(env, k, j)
    best_val, best_plan = solver.solve()

    return SCN_P[k, j] * best_val


def cost_one_particle(prop: np.ndarray) -> float:
    """对单粒子 (α, β) 计算 25 场景期望收益 (负号→成本)"""
    alpha, beta = prop
    print(f"[粒子] α={alpha:.4f}, β={beta:.4f}", flush=True)
    if alpha < 0 or beta < 0 or alpha + beta > 1:
        return LARGE_PN

    scene_args = [(k, j, alpha, beta) for k in range(5) for j in range(5)]
    with parallel_backend("loky", inner_max_num_threads=1):
        scene_vals = Parallel(n_jobs=-1, batch_size=1)(
            delayed(_solve_one_scene)(k, j, alpha, beta)
            for k, j, alpha, beta in scene_args
        )

    total_val = sum(scene_vals)
    return -total_val


# ------------------------------------------------------------
#  5. PSO 优化（保持原有逻辑）
# ------------------------------------------------------------


def main():
    print("[Numba迭代DP] 开始搜索 α / β …", flush=True)
    optimizer = ps.single.GlobalBestPSO(
        n_particles=20,
        dimensions=2,
        options={"c1": 1.2, "c2": 1.2, "w": 0.6},
        bounds=(np.zeros(2), np.ones(2)),
    )

    def objective_fn(X: np.ndarray) -> np.ndarray:
        return np.array([cost_one_particle(x) for x in X])

    with parallel_backend("loky", inner_max_num_threads=1):
        best_cost, best_pos = optimizer.optimize(objective_fn, iters=30, verbose=False)
    a, b = best_pos
    print(f"[完成] α={a:.4f},β={b:.4f},收益≈{-best_cost:.2f}")
    print(f"迭代DP调用次数: {hit['iterative_calls']}")
    plt.figure()
    plt.plot(optimizer.cost_history)
    plt.xlabel("iter")
    plt.ylabel("cost")
    plt.savefig("pso_convergence_iterative.png")
    print("已保存收敛曲线")


def instant_revenue(
    t: int,
    power: float,
    con: np.ndarray,
    da: np.ndarray,
    p_da: np.ndarray,
    p_rt: np.ndarray,
    wind: np.ndarray,
) -> float:
    """单小时收益函数（保持与原版兼容）"""
    return instant_revenue_numba(
        t, power, con[t], da[t], p_da[t], p_rt[t], wind[t], BASE_PRICE[t], REF_PRICE[t]
    )


if __name__ == "__main__":
    from multiprocessing import freeze_support

    freeze_support()
    main()